import numpy as np
import cv2 as cv
from keras.models import load_model
from PIL import Image
import pygame
import sys

# 加载模型
model = load_model('my_model.keras')

def pre_pic(picName):
    img = Image.open(picName)
    reIm = img.resize((100, 100), Image.Resampling.LANCZOS)
    im_arr = np.array(reIm.convert("L"))
    return im_arr

# 初始化pygame
pygame.init()
size = (1280, 480)
screen = pygame.display.set_mode(size)
pygame.display.set_caption("剪刀石头布 - 对战结果")

# 定义颜色
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

# AI选择获胜的逻辑
def ai_winning_choice(player_choice):
    if player_choice == 0:  # 剪刀
        return 1  # 石头
    elif player_choice == 1:  # 石头
        return 2  # 布
    else:  # 布
        return 0  # 剪刀

# 显示图像
def display_choice(player, ai):
    screen.fill(WHITE)
    font = pygame.font.Font(None, 50)
    text = font.render(f"Player: {player}, AI: {ai}", True, BLACK)
    screen.blit(text, (50, 30))
    pygame.display.flip()

cap = cv.VideoCapture(0)
filename = '00.jpg'

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            cap.release()
            cv.destroyAllWindows()
            pygame.quit()
            sys.exit()

    success, frame = cap.read()
    if not success:
        break
    new_img = cv.flip(frame, 1)  # 镜像翻转
    roi = new_img.copy()

    # 肤色检测
    ycrcb = cv.cvtColor(roi, cv.COLOR_BGR2YCrCb)
    cr = cv.split(ycrcb)[1]
    cr1 = cv.GaussianBlur(cr, (5, 5), 0)
    _, skin1 = cv.threshold(cr1, 0, 255, cv.THRESH_BINARY + cv.THRESH_OTSU)

    # 查找轮廓
    contours, _ = cv.findContours(skin1, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    detected = False
    choice_text = ["jiandao", "shitou", "bu"]
    player_choice = -1

    for item in contours:
        (x, y, w, h) = cv.boundingRect(item)
        if w >= 50 and h >= 50:
            img = skin1[y:y+h, x:x+w]
            cv.imwrite(filename, img)
            img = pre_pic(filename)
            img = img.reshape(1, 1, 100, 100).astype('float32')
            img /= 255.0
            predictions = np.argmax(model.predict(img), axis=1)
            flag = model.predict(img)
            if max(flag[0]) > 0.6:
                player_choice = predictions[0]
                ai_choice = ai_winning_choice(player_choice)
                print(f"Player: {choice_text[player_choice]}, AI: {choice_text[ai_choice]}")
                display_choice(choice_text[player_choice], choice_text[ai_choice])

                # 绘制识别框和文字
                cv.rectangle(new_img, (x, y), (x + w, y + h), (255, 0, 0), 2)
                cv.putText(new_img, choice_text[player_choice], (x, y - 10),
                           cv.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
                detected = True

    # 如果没有检测到手势，则清空pygame显示
    if not detected:
        screen.fill(WHITE)
        pygame.display.flip()

    # 将三个图像都转为三通道图像以便拼接
    skin_color = cv.cvtColor(skin1, cv.COLOR_GRAY2BGR)

    gray_roi = cv.cvtColor(roi, cv.COLOR_BGR2GRAY)
    gray_3ch = cv.cvtColor(gray_roi, cv.COLOR_GRAY2BGR)  # 灰度图转为三通道

    new_img_with_box = new_img  # 已经是三通道

    # 拼接三个视图：皮肤检测、原始灰度图、带识别框的图像
    combined_img = np.hstack((skin_color, gray_3ch, new_img_with_box))

    # 调整尺寸以适配屏幕（可选）
    combined_img = cv.resize(combined_img, (1280, 480))

    cv.imshow('剪刀石头布 - 视频流综合', combined_img)

    Key = cv.waitKey(1)
    if Key == ord(' '):  # 按下空格键退出
        break

cap.release()
cv.destroyAllWindows()
pygame.quit()
sys.exit()